import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import cv2
from insightface.app import FaceAnalysis
from pymilvus import MilvusClient

video_capture = cv2.VideoCapture(0)
app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])  # 优先GPU加速
app.prepare(ctx_id=0, det_size=(640, 640))  # 第一个使用GPU， 检测分辨率

cli = MilvusClient(
    uri='http://39.104.78.210:19530',
    token='root:Milvus',
) # 建立连接

while True:
    _, frame = video_capture.read()
    faces = app.get(frame)
    for face in faces:
        embedding = face.embedding
        res = cli.search(
            collection_name="faceD",
            anns_field="vector",
            data=[embedding],
            limit=1,
            output_fields=["name"],
            search_params={"metric_type": "COSINE"}
        )  # ANN(KNN) 查找最相近的元素
        if res[0][0]['distance'] > 0.5:
            print(res)
    cv2.imshow('Video', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

video_capture.release()
cv2.destroyAllWindows()